"""
基于 corotbeam.py 的 DTU10MW 叶片结构分析示例

此示例演示如何使用 elements_created_from_file 类从结构数据文件
创建共旋转梁单元，并进行基本的分析。
"""

import os
import sys
import numpy as np
from matplotlib import pyplot as plt

# 导入 corotbeam 模块
# 假设 corotbeam.py 在 test_all 目录下，或者已安装到包中
##sys.path.append(os.path.join(os.path.dirname(__file__), '../../'))
from casetoolbox.casestab.corotbeam import elements_created_from_file, equilibrium_beam_element, corotbeam_element_kinematics

# 设置工作目录和数据路径
work_folder = os.path.join(os.path.dirname(__file__))
data_file = os.path.join(work_folder, 'stru_6x6.dat')  # 或使用 stru_ISO.dat

# 定义节点分布（从 DTU10MW_6x6.json 中的 znode）
znode = np.array([
    0.0000, 3.0000, 6.0000, 7.0000, 8.7005, 10.4020, 12.2046, 13.2065, 
    15.0100, 18.2151, 21.4178, 24.6189, 27.8193, 31.0194, 34.2197, 
    40.2204, 46.6217, 53.0232, 59.4245, 65.8255, 72.2261, 79.0266, 
    80.5267, 82.0271, 83.5274, 85.0277, 86.366
])

# 参数设置
setno = 1          # 数据集编号
subsetno = 1       # 子集编号
nintp = 50         # 每个单元内的插值点数
norder = 0         # 单元阶数（0 = 线性单元）

print("=" * 60)
print("DTU10MW 叶片结构分析示例")
print("=" * 60)
print(f"数据文件: {data_file}")
print(f"节点数: {len(znode)}")
print(f"单元数: {len(znode) - 1}")
print(f"插值点数: {nintp}")
print(f"单元阶数: {norder}")
print("=" * 60)

# 创建元素
print("\n正在读取结构数据并创建单元...")
elements = elements_created_from_file(
    fname=data_file,
    znode=znode,
    setno=setno,
    subsetno=subsetno,
    nintp=nintp,
    norder=norder
)

print(f"成功创建 {elements.nelem} 个单元")
print(f"单元位置范围: {elements.z[0]:.2f} m 到 {elements.z[-1]:.2f} m")

# 分析单元属性
print("\n" + "=" * 60)
print("单元属性分析")
print("=" * 60)

# 计算每个单元的长度和质量
unit_lengths = elements.l
total_length = np.sum(unit_lengths)

# 从结构数据中提取质量信息
mass_per_length = elements.stset[:, 4]  # 第5列是质量/长度
z_coords = elements.stset[:, 0]          # 第1列是z坐标

print(f"\n叶片总长度: {total_length:.2f} m")
print(f"单元长度范围: {unit_lengths.min():.2f} m 到 {unit_lengths.max():.2f} m")
print(f"平均单元长度: {unit_lengths.mean():.2f} m")

# 计算总质量（通过积分）
# 使用梯形积分
total_mass = np.trapz(mass_per_length, z_coords)
print(f"叶片总质量: {total_mass/1000:.2f} t (约 {total_mass:.0f} kg)")

# 分析单元柔度矩阵
print("\n" + "=" * 60)
print("单元柔度矩阵分析")
print("=" * 60)

# 显示第一个和最后一个单元的柔度矩阵
print("\n第一个单元 (根部) 的柔度矩阵 C[0]:")
print("轴向柔度 C11 [1/N]:")
print(f"  {elements.Cs[0][0][0,0]:.6e}")

print("\n最后一个单元 (叶尖) 的柔度矩阵 C[0]:")
last_elem_idx = elements.nelem - 1
print(f"  {elements.Cs[last_elem_idx][0][0,0]:.6e}")

# 分析单元惯性参数
print("\n" + "=" * 60)
print("单元惯性参数分析")
print("=" * 60)

print("\n第一个单元的惯性参数:")
first_iner = elements.iner_pars[0]
print(f"  质量密度 m: {first_iner[0,0]:.2f} kg/m")
print(f"  质量矩 mx_cg: {first_iner[0,1]:.2f} kg·m/m")
print(f"  质量矩 my_cg: {first_iner[0,2]:.2f} kg·m/m")

print("\n最后一个单元的惯性参数:")
last_iner = elements.iner_pars[last_elem_idx]
print(f"  质量密度 m: {last_iner[0,0]:.2f} kg/m")
print(f"  质量矩 mx_cg: {last_iner[0,1]:.2f} kg·m/m")
print(f"  质量矩 my_cg: {last_iner[0,2]:.2f} kg·m/m")

# 可视化
print("\n" + "=" * 60)
print("生成可视化图表...")
print("=" * 60)

fig, axes = plt.subplots(2, 2, figsize=(12, 10))

# 1. 单元长度分布
ax1 = axes[0, 0]
ax1.plot(elements.z, unit_lengths, 'b-o', markersize=4)
ax1.set_xlabel('沿叶片轴向位置 z [m]')
ax1.set_ylabel('单元长度 [m]')
ax1.set_title('单元长度分布')
ax1.grid(True)

# 2. 质量分布
ax2 = axes[0, 1]
ax2.plot(z_coords, mass_per_length, 'r-', linewidth=2)
ax2.set_xlabel('沿叶片轴向位置 z [m]')
ax2.set_ylabel('单位长度质量 [kg/m]')
ax2.set_title('叶片质量分布')
ax2.grid(True)

# 3. 单元位置
ax3 = axes[1, 0]
ax3.plot(elements.pos[0, :], elements.pos[1, :], 'g-o', markersize=4)
ax3.set_xlabel('x 位置 [m]')
ax3.set_ylabel('y 位置 [m]')
ax3.set_title('叶片参考曲线 (俯视图)')
ax3.axis('equal')
ax3.grid(True)

# 4. 柔度矩阵元素沿叶片的变化
ax4 = axes[1, 1]
# 提取轴向柔度 C11
C11_values = []
z_interp = []
for ielem in range(elements.nelem):
    z_interp.append(0.5 * (elements.z[ielem] if ielem > 0 else 0) + 
                    0.5 * elements.z[ielem] if ielem < len(elements.z) else elements.z[-1])
    C11_values.append(elements.Cs[ielem][0][0, 0])

ax4.semilogy(z_interp, C11_values, 'b-o', markersize=4)
ax4.set_xlabel('沿叶片轴向位置 z [m]')
ax4.set_ylabel('轴向柔度 C11 [1/N]')
ax4.set_title('轴向柔度沿叶片的变化')
ax4.grid(True)

plt.tight_layout()

# 保存图片
output_dir = os.path.join(work_folder, 'corotbeam_example_results')
os.makedirs(output_dir, exist_ok=True)
output_file = os.path.join(output_dir, 'corotbeam_analysis.png')
plt.savefig(output_file, dpi=150)
print(f"\n图表已保存到: {output_file}")

# 创建平衡梁单元示例（仅用于演示）
print("\n" + "=" * 60)
print("创建平衡梁单元示例")
print("=" * 60)

# 取第一个单元创建平衡梁单元
ielem = 0
elem_model = equilibrium_beam_element(
    l=elements.l[ielem],
    C=elements.Cs[ielem],
    iner_pars=elements.iner_pars[ielem]
)

print(f"\n单元 {ielem} 的属性:")
elem_model.print_element_properties()

# 创建单元运动学对象
elem_state = corotbeam_element_kinematics(
    l0=elements.l[ielem],
    r0=elements.r0s[ielem],
    E0=elements.E0s[ielem],
    norder=norder
)

print(f"\n单元 {ielem} 的初始长度: {elem_state.l0:.4f} m")
print(f"单元 {ielem} 的初始位置向量:")
print(f"  起点: [{elements.r0s[ielem][0]:.4f}, {elements.r0s[ielem][1]:.4f}, {elements.r0s[ielem][2]:.4f}]")
print(f"  终点: [{elements.r0s[ielem][3]:.4f}, {elements.r0s[ielem][4]:.4f}, {elements.r0s[ielem][5]:.4f}]")

print("\n" + "=" * 60)
print("示例完成!")
print("=" * 60)